// 这里就是 https://juejin.cn/book/7347579913702293567/section/7355761778120196105 这里的动手环节啦
// 相比于rag-demo 这里主要是增加了 memory 的相关 feature

import { FaissStore } from "@langchain/community/vectorstores/faiss";
import { useEmbeddings } from "../rag-demo/utils";
import { config } from "dotenv";
import path from "path";
import { fileURLToPath } from "url";
import {
  RunnablePassthrough,
  RunnableSequence,
  RunnableWithMessageHistory,
} from "@langchain/core/runnables";
import { useTemplatePrompt } from "./Template";
import { useModel } from "../model";
import { StringOutputParser } from "@langchain/core/output_parsers";
import { useRephraseChain } from "./RephraseChain";
import { JSONChatHistory } from "./history";
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

/**
 * 举个例子，我们用户在对话中，常常有：他，那个， 这个，它  ... 等等指代词，没有上下文联系的能力的话，llm 是无法获悉用户指代的对象的
 * 为了获取更好的提问质量，我们需要改下提问，使得用户的提问成为一个独立的问题
 * 怎么改写呢，用什么改写呢？ 依旧是用 llm，我们来定义一个 chain 来完成这个任务, 即rephrase chain
 */

async function loadVectorStore(dir: string) {
  // 向量化编码
  const embeddings = useEmbeddings();
  const vectorStore = await FaissStore.load(dir, embeddings);
  return vectorStore;
}

async function getRagChain() {
  await config();
  const dir = path.resolve(__dirname, "..", "..", "data", "db", "qiu");
  const vectorStore = await loadVectorStore(dir);
  const retriever = vectorStore.asRetriever(2);
  const convertDocsToString = (documents: any[]): string => {
    return documents.map((document) => document.pageContent).join("\n");
  };
  const contextRetrieverChain = RunnableSequence.from([
    (input) => input.standalone_question,
    retriever,
    convertDocsToString,
  ]);
  const prompt = useTemplatePrompt();
  const model = useModel({});
  const rephraseChain = await useRephraseChain();
  const ragChain = RunnableSequence.from([
    // 用 rephraseChain 来改写用户问题， 并把改写后的问题放到变量standalone_question中
    RunnablePassthrough.assign({
      standalone_question: rephraseChain,
    }),
    // 用 contextRetrieverChain 来获取上下文， 并把上下文放到变量context中
    RunnablePassthrough.assign({
      context: contextRetrieverChain,
    }),
    // 这个 prompt需要输入：context, standalone_question, history
    prompt,
    model,
    new StringOutputParser(),
  ]);
  // 那么我们可以看到，ragChain 其实还没有 history 信息，那我们这里用RunnableWithMessageHistory取管理
  const ragChainWithHistory = new RunnableWithMessageHistory({
    runnable: ragChain,
    // 一个用于存储 history 的本地目录
    getMessageHistory: (_sessionId: any) =>
      new JSONChatHistory({
        sessionId: _sessionId,
        dir: path.resolve(__dirname, "..", "..", "data", "db", "chat-history"),
      }),
    historyMessagesKey: "history",
    inputMessagesKey: "question",
  });
  return ragChainWithHistory;
  //   const res = await ragChainWithHistory.invoke(
  //     {
  //       question: "什么是球状闪电",
  //     },
  //     {
  //       configurable: {
  //         sessionId: "user_session_001", // 用于标识不同用户或不同会话
  //       },
  //     }
  //   );
}

async function main() {
  const ragChainWithHistory = await getRagChain();
  const res = await ragChainWithHistory.invoke(
    {
      question: "这个现象在文中有什么故事",
    },
    {
      configurable: {
        sessionId: "user_session_001", // 用于标识不同用户或不同会话
      },
    }
  );
  console.log(res);
}
// main(); // 注释掉，避免在 server.ts 导入时自动执行

export { getRagChain };